//
// Created by swx on 23-6-4.
//
#include <iostream>
#include <thread>
#include <vector>
#include <random>
#include "clerk.h"
#include <atomic>

atomic<int> totalCount(0); // 统计总操作次数
// atomic<int> successCount(0);

// 生成从 "AA" 到 "ZZ" 的键
std::vector<std::string> generateKeys() {
    std::vector<std::string> keys;
    // for (char i = 'A'; i <= 'Z'; ++i) {
    //     for (char j = 'A'; j <= 'Z'; ++j) {
    //         keys.emplace_back(std::string(1, i) + j);
    //     }
    // }
    for (char i = 'A'; i <= 'Z'; ++i) {
        keys.emplace_back(std::string(1, i));
    }
    return keys;
}

void stressTest(Clerk& client, const std::vector<std::string>& keys, int operationsPerThread) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> keyDist(0, keys.size() - 1);
    std::uniform_real_distribution<> opDist(0.0, 1.0);

    for (int i = 0; i < operationsPerThread; ++i) {
        std::string key = keys[keyDist(gen)];
        if (opDist(gen) < 0.2) { // 插入操作的概率为 0.2
            std::string value = std::to_string(gen());
            client.Put(key, value);
        }
        else { // 读取操作
            std::string value = client.Get(key);
            // std::cout << "Thread " << std::this_thread::get_id() << ": Get key {" << key << "}, value {" << value << "}" << std::endl;
        }
        totalCount++;
    }
}

void stressTest2(Clerk& client, const std::vector<std::string>& keys, int operationsPerThread) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> keyDist(0, keys.size() - 1);
    std::uniform_real_distribution<> opDist(0.0, 1.0);
    std::string key = "";
    for (int i = 0; i < operationsPerThread; ++i) {
        if (key == "" || opDist(gen) < 0.0) { // 插入操作的概率为 0.2
            key = keys[keyDist(gen)];
            std::string value = std::to_string(gen());
            client.Put(key, value);
        }
        else { // 读取操作
            std::string value = client.Get(key);
            // std::cout << "Thread " << std::this_thread::get_id() << ": Get key {" << key << "}, value {" << value << "}" << std::endl;
        }
        totalCount++;
        if (totalCount % 10 == 0) {
            std::cout << "Total operations: " << totalCount.load() << std::endl;
        }
    }
}

int main() {
    Clerk client;
    client.Init("test.conf");

    auto start = std::chrono::high_resolution_clock::now();

    const int threadCount = 1; // 线程数
    const int operationsPerThread = 1000; // 每个线程的操作次数
    std::vector<std::string> keys = generateKeys(); // 生成键范围 "A" 到 "Z"

    std::vector<std::thread> threads;
    for (int i = 0; i < threadCount; ++i) {
        threads.emplace_back(stressTest2, std::ref(client), std::ref(keys), operationsPerThread);
    }

    for (auto& t : threads) {
        t.join();
    }

    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
    std::cout << "Test completed in " << duration.count() << " us" << std::endl;
    std::cout << "QPS: " << (1000000L * threadCount * operationsPerThread) / duration.count() << " ops/s" << std::endl;

    return 0;
}